package main

import (
	"encoding/json"
	"fmt"
	"math/rand"
	"os"
	"time"
)

// 初始化随机种子
func init() {
	rand.Seed(time.Now().UnixNano())
}

// 持久化ID映射表结构
type IDMapping struct {
	CharToID map[string]int `json:"char_to_id"`
	IDToChar map[int]string `json:"id_to_char"`
	NextID   int            `json:"next_id"`
}

// Load 从文件加载模型
func LoadTransformerModel(filePath string) (*TransformerLanguageModel, error) {
	// 读取文件
	jsonData, err := os.ReadFile(filePath)
	if err != nil {
		return nil, fmt.Errorf("读取模型文件失败: %v", err)
	}

	// 反序列化JSON
	var modelData struct {
		Vocab        map[string]int     `json:"vocab"`
		ReverseVocab map[int]string     `json:"reverse_vocab"`
		VocabSize    int                `json:"vocab_size"`
		Embedding    [][]float32        `json:"embedding"`
		PosEncodings [][]float32        `json:"pos_encodings"`
		Layers       []TransformerLayer `json:"layers"`
		OutputProj   [][]float32        `json:"output_proj"`
		DModel       int                `json:"d_model"`
		NLayer       int                `json:"n_layer"`
		MaxLen       int                `json:"max_len"`
	}

	if err := json.Unmarshal(jsonData, &modelData); err != nil {
		return nil, fmt.Errorf("反序列化模型失败: %v", err)
	}

	// 创建模型
	model := &TransformerLanguageModel{
		Vocab:           modelData.Vocab,
		ReverseVocab:    modelData.ReverseVocab,
		VocabSize:       modelData.VocabSize,
		Embedding:       modelData.Embedding,
		PosEncodings:    modelData.PosEncodings,
		Layers:          modelData.Layers,
		OutputProj:      modelData.OutputProj,
		DModel:          modelData.DModel,
		NLayer:          modelData.NLayer,
		MaxLen:          modelData.MaxLen,
		DropoutRate:     0.3, // 默认值
		TrendWindowSize: 5,   // 默认值
	}

	return model, nil
}

// Save 保存模型到文件
func (lm *TransformerLanguageModel) SaveTransformerModel(filePath string) error {
	// 创建模型保存数据结构
	modelData := struct {
		Vocab        map[string]int     `json:"vocab"`
		ReverseVocab map[int]string     `json:"reverse_vocab"`
		VocabSize    int                `json:"vocab_size"`
		Embedding    [][]float32        `json:"embedding"`
		PosEncodings [][]float32        `json:"pos_encodings"`
		Layers       []TransformerLayer `json:"layers"`
		OutputProj   [][]float32        `json:"output_proj"`
		DModel       int                `json:"d_model"`
		NLayer       int                `json:"n_layer"`
		MaxLen       int                `json:"max_len"`
	}{}

	// 填充数据
	modelData.Vocab = lm.Vocab
	modelData.ReverseVocab = lm.ReverseVocab
	modelData.VocabSize = lm.VocabSize
	modelData.Embedding = lm.Embedding
	modelData.PosEncodings = lm.PosEncodings
	modelData.Layers = lm.Layers
	modelData.OutputProj = lm.OutputProj
	modelData.DModel = lm.DModel
	modelData.NLayer = lm.NLayer
	modelData.MaxLen = lm.MaxLen

	// 序列化到JSON
	jsonData, err := json.MarshalIndent(modelData, "", "  ")
	if err != nil {
		return fmt.Errorf("序列化模型失败: %v", err)
	}

	// 写入文件
	if err := os.WriteFile(filePath, jsonData, os.ModePerm); err != nil {
		return fmt.Errorf("写入模型文件失败: %v", err)
	}

	return nil
}
